import numpy as np
import numpy.random as npr
import copy, string, itertools, random, bisect

#global Nc, Nl, Na

Nc = 5    # number of chromosomes
Nl = 4    # number of gene loci per chromosomes
Na = 10   # number of alleles per locus
F  = 10   # fecundity = number of offspring per female

transmission_rate = 0.9   # transmission rate of the male-killer
lethality = 0.8  # killing efficiency
reallocation_efficiency = 1.

Nl_tot = Nc*Nl    # total number of loci
cnames = chromosome_names = ["Chr{0}{1}".format(i,s) for s in ['A','B'] for i in range(1,Nc+1)]
gene_names_flat = ["loc{0}".format(i) for i in range(1,Nl_tot+1)]
gene_names = [gene_names_flat[i*Nl:(i+1)*Nl] for i in range(Nc)]
anames = allele_names = [["a{0}_{1}".format(i,j) for j in range(1,Na+1)] for i in range(1,Nl_tot+1)]
allele_names_flat = list(itertools.chain.from_iterable(allele_names))

boolval = {'uninfected': 0, 'infected': 1}

phenotypic_value = dict()
for allele in allele_names_flat:
    phenotypic_value[allele] = npr.random()

max_phenotypic_value = 2 * Nc * Nl

def create_random_chromosomes():
    chromosomes = [[] for i in range(2*Nc)]     # empty diploid chromosome set
    for i in range(2):
        count = 0
        for j in range(i*Nc,(i+1)*Nc):
            for k in range(Nl):
                chromosomes[j].append(random.choice(allele_names[count]))
                count += 1
    return chromosomes
    
class WeightedRandomGenerator(object):
    def __init__(self, weights):
        self.totals = []
        running_total = 0

        for w in weights:
            running_total += w
            self.totals.append(running_total)

    def next(self):
        rnd = random.random() * self.totals[-1]
        return bisect.bisect_right(self.totals, rnd)

    def __call__(self):
        return self.next()

class WeightedRandomGeneratorWithoutReplacement(object):
    def __init__(self, weights):
        self.totals = []
        self.selected = []
        running_total = 0

        for w in weights:
            running_total += w
            self.totals.append(running_total)

    def next(self):
        if len(self.totals) == len(self.selected):
            raise ValueError, "generator is empty"
        rnd = random.random() * self.totals[-1]
        i = bisect.bisect_right(self.totals, rnd)
        if i not in self.selected:
            self.selected.append(i)
            return i
        else:
            return self.next()
        
    def __call__(self):
        return self.next()


class Individual(object):
    def __init__(self, chromosomes, infection_status, sex=None, fitness=1.):
        self.chromosomes = chromosomes
        self.chromosome_set1 = chromosomes[:Nc]
        self.chromosome_set2 = chromosomes[Nc:]
        self.infection_status = infection_status
        self.infected = boolval[infection_status]
        self.sex = sex
        self.fitness = fitness
        self.phenotype = self.get_phenotypic_value()
        
    def __str__(self):
        return self.show(verbose=False, direct_call=False)
    
    def show(self, verbose=True, direct_call=True):
        if not verbose:
            s = """{}: {}
    phenotype: {}
    fitness: {}
    """.format( self.sex, self.infection_status, self.phenotype, self.fitness )
        else:
            s = """{}: {}
    phenotype: {}
    fitness: {}
    genotype:
    {}
    allelic phenotypic values:
    {}
    """.format( self.sex, self.infection_status, self.phenotype, self.fitness, self.chromosomes2string(), self.phenotypes2string() )
        if direct_call:
            print s
        else:
            return s
    
    def chromosomes2string(self):
        s = ''
        for i in range(Nc):
            s += "{0:>7}  {1:<7} | ".format(cnames[i],cnames[Nc+i])
        s = s.rstrip(' |')
        l = len(s)
        s += '\n' + '~'*l + '\n'
        for i in range(Nl):
            for j in range(Nc):
                s += "{0:>7}  {1:<7} | ".format(self.chromosomes[j][i],self.chromosomes[Nc+j][i])
            s = s.rstrip(' |')
            s += '\n'
        return s
    
    def phenotypes2string(self):
        s = ''
        for i in range(Nc):
            s += "{0:>7}  {1:<7} | ".format(cnames[i],cnames[Nc+i])
        s = s.rstrip(' |')
        l = len(s)
        s += '\n' + '~'*l + '\n'
        for i in range(Nl):
            for j in range(Nc):
                s += "{0:>7.3f}  {1:<7.3f} | ".format(phenotypic_value[self.chromosomes[j][i]],phenotypic_value[self.chromosomes[Nc+j][i]])
            s = s.rstrip(' |')
            s += '\n'
        return s
    
    def get_phenotypic_value(self):
        value = 0.
        for ch in self.chromosomes:
            for a in ch:
                value += phenotypic_value[a]
        return value/max_phenotypic_value
    
    def produce_gametes(self):
        gametes = [[] for i in range(4)]
        for i in range(Nc):
            for j in range(4):
                h = recombine(self.chromosome_set1[i], self.chromosome_set2[i])
                random.shuffle(h)
                gametes[j].append( h[j] )
        return gametes
    
    def update_fitness(self, factor):
        self.fitness *= factor
        

class Female(Individual):
    def __init__(self, chromosomes, infection_status, fitness=1.):
        Individual.__init__(self, chromosomes, infection_status, sex='female', fitness=fitness)
        self.sex_ratio = self.phenotype
        self.fecundity = F
    
    def produce_brood_with(self, male):
        daughters = []
        sons = []
        for i in range(self.fecundity):
            # offspring chromosomes:
            egg = random.choice( self.produce_gametes() )
            sperm = random.choice( male.produce_gametes() )
            # offspring infection status:
            offspring_status = 'uninfected'
            if self.infected:
                if npr.random() < transmission_rate:               # random transmission of infection
                    offspring_status = 'infected'
            # offspring sex:
            if npr.random() < self.sex_ratio:
                sons.append( Male(egg+sperm, offspring_status) )
            else:
                daughters.append( Female(egg+sperm, offspring_status) )
        if self.infected:
            sons,daughters = self.male_killing(sons,daughters)
        return sons, daughters
    
    def male_killing(self, sons, daughters):
        sons[:] = (s for s in sons if s.infected and not npr.random()<lethality)
        survivers = len(sons+daughters)
        body_count = self.fecundity - survivers
        # resource reallocation:
        fitness_compensation = 1. + reallocation_efficiency*body_count/survivers
        for offspring in sons+daughters:
            offspring.update_fitness(fitness_compensation)
        return sons,daughters
    
class Male(Individual):
    def __init__(self, chromosomes, infection_status, fitness=1.):
        Individual.__init__(self, chromosomes, infection_status, sex='male', fitness=fitness)
        self.number_of_matings = 0

class Population(object):
    def __init__(self, females, males):
        self.females = females
        self.males = males
        self.offspring = []
        self.generation = 0
        self.secondary_male_mating_pool = []
        
    def next_generation(self):
        n_females = len(self.females)
        n_males = len(self.males)
        for female in self.females:
            pass
        
        

def recombine(c1,c2):
    """Recombine two homologous chromosomes """
    if len(c1) != len(c2):
        raise ValueError, "Chromosomes must have the same length in order to properly recombine."

    # meiosis I
    # S phase
    h1 = copy.copy(c1)
    h2 = copy.copy(c1)
    h3 = copy.copy(c2)
    h4 = copy.copy(c2)
    
    # prophase I
    pos = npr.random_integers(1, Nl-1, 1)    # randomly choose position of chrossing of the two homologous chromosomes
    h2[:pos] = c2[:pos]
    h3[:pos] = c1[:pos]
    
    # meiosis II
    return [h1, h2, h3, h4]
    
if __name__ == '__main__':
    npr.seed()
    random.seed()
    c1 = create_random_chromosomes()
    c2 = create_random_chromosomes()
    
    f = Female(c1, 'infected')
    m = Male(c2, 'uninfected')
    print f
    print m
    
    print "~~~~~~~~~~~~~~~~~~ offspring ~~~~~~~~~~~~~~~~~~~"
    sons,daughters = f.produce_brood_with(m)
    
    for o in sons+daughters:
        print o

    
    
